package org.apache.flink.runtime.resourceconsumer;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.apache.flink.runtime.allocatedeploy.execution.ExecutionID;
import org.apache.flink.runtime.clusterframework.types.ResourceID;

/** Default {@link ExecutionDeploymentTracker} implementation. */
public class DefaultExecutionDeploymentTracker implements ExecutionDeploymentTracker {

    private final Set<ExecutionID> pendingDeployments = new HashSet<>();
    private final Map<ResourceID, Set<ExecutionID>> executionsByHost = new HashMap<>();
    private final Map<ExecutionID, ResourceID> hostByExecution = new HashMap<>();

    @Override
    public void startTrackingPendingDeploymentOf(
    		ExecutionID executionAttemptId, ResourceID host) {
        pendingDeployments.add(executionAttemptId);
        hostByExecution.put(executionAttemptId, host);
        executionsByHost.computeIfAbsent(host, ignored -> new HashSet<>()).add(executionAttemptId);
    }

    @Override
    public void completeDeploymentOf(ExecutionID executionId) {
        pendingDeployments.remove(executionId);
    }

    @Override
    public void stopTrackingDeploymentOf(ExecutionID executionId) {
        pendingDeployments.remove(executionId);
        ResourceID host = hostByExecution.remove(executionId);
        if (host != null) {
            executionsByHost.computeIfPresent(
                    host,
                    (resourceID, executionAttemptIds) -> {
                        executionAttemptIds.remove(executionId);

                        return executionAttemptIds.isEmpty() ? null : executionAttemptIds;
                    });
        }
    }

    @Override
    public Map<ExecutionID, ExecutionDeploymentState> getExecutionsOn(ResourceID host) {
        return executionsByHost.getOrDefault(host, Collections.emptySet()).stream()
                .collect(
                        Collectors.toMap(
                                x -> x,
                                x ->
                                        pendingDeployments.contains(x)
                                                ? ExecutionDeploymentState.PENDING
                                                : ExecutionDeploymentState.DEPLOYED));
    }
}
